from bloompy import BloomFilter as BloomFilter1, CountingBloomFilter as CountingBloomFilter1
import mmh3
import socket
from config import Config
import dpkt
import datetime

def inet_to_str(inet):
    """Convert inet object to a string
        Args:
            inet (inet struct): inet network address
        Returns:
            str: Printable/readable IP address
    """
    # First try ipv4 and then ipv6
    try:
        return socket.inet_ntop(socket.AF_INET, inet)
    except ValueError:
        return socket.inet_ntop(socket.AF_INET6, inet)

class BloomFilter(BloomFilter1):
    def reset(self):
        self.__init__(self.error_rate, self._c, self._b)


class CountingBloomFilter(CountingBloomFilter1):
    def add(self,element):
        '''
        query the element status in the filter and add it into the filter
        '''
        if self.count >= self.capacity:
            raise IndexError('BloomFilter is at capacity.')
        # if element in self:
        #     return True
        _element = self._to_str(element)
        i = 0
        for _ in range(self.hash_num):
            hashed_value = mmh3.hash(_element,
                                        self.seeds[_]) % self.bit_num
            raw_value = self._get_bit_value(hashed_value)
            if raw_value > 0:
                i+=1
            self._set_bit_value(hashed_value,raw_value+1)
        if i == self.hash_num:
            return True
        else:
            self.count += 1
            return False

    def get_counter_num(self, element):
        """Get the count of the given element.
        """
        nums = []
        _element = self._to_str(element)
        for _ in range(self.hash_num):
            hashed_value = mmh3.hash(_element,
                                        self.seeds[_])%self.bit_num
            nums.append(self._get_bit_value(hashed_value))
        return min(nums)



def print_quintuple(file_path, n=5):
    """Print out the first n matched IP Packets' information from given pacp file,
    where the information includes but not limited to:
    - timestamp
    - source IP address
    - destination IP address
    - sorcec port number
    - destination port number
    - portocol type
    ...

    Args:
        file_path: str, ".pacp" format.
        n: int, the number of packet that will be parsed. 
            If n is less than 0, then all packets will be parsed.

    """
    ctr = 0 # counter
       
    with open(Config.PCAP_FILE_PATH, mode='rb') as f_pcap:

        # Read pcap file
        pcap = dpkt.pcap.Reader(f_pcap)

        for ts, buf in pcap:

            if ctr == n: break
            ctr += 1

            # Print out the timestamp in UTC
            print("timestamp: ", str(datetime.datetime.utcfromtimestamp(ts)))

            # Unpack the IP Packet
            ip_packet = dpkt.ip.IP(buf)

            # Make sure the 'buf' is IP Packet
            if not isinstance(ip_packet, dpkt.ip.IP):
                print('Non IP Packet type not supported %s\n' % buf.__class__.__name__)
                continue
            
            # Print out IP Packet information
            print("src IP: ", inet_to_str(ip_packet.src))
            print("dst IP: ", inet_to_str(ip_packet.dst))

            # Make sure ip_packet.data is TCP Packet
            tcp_packet = ip_packet.data

            if not isinstance(tcp_packet, dpkt.tcp.TCP):
                print('Non TCP Packet type not supported %s\n' % tcp_packet.__class__.__name__)
                continue         

            # Print out TCP Packet information
            print("src port: ", tcp_packet.sport)
            print("dst port: ", tcp_packet.dport)

def test_BloomFilter():
    bf = BloomFilter(Config.error_rate, Config.element_num)
    print(1 in bf)
    bf.add(1)
    print(1 in bf)
    bf.reset()
    print(1 in bf)

if __name__ == '__main__':
    test_BloomFilter()